import numpy as np
import os
import operator

class_2_label_lower = [
    [240, 0, 0], [251, 255, 3], [44, 251, 255], [241, 47, 255], [125, 255, 155], [26, 125, 255], [255, 234, 157],
    [206, 129, 212], [45, 135, 66], [185, 207, 45], [69, 147, 207], [207, 72, 104], [4, 207, 4], [35, 1, 207]
]

class_2_label_upper = [
    [170, 255, 127], [170, 255, 255], [255, 255, 0], [255, 170, 0], [170, 170, 255], [0, 170, 255], [85, 170, 0],
    [255, 85, 255], [255, 85, 127], [85, 170, 127], [255, 85, 0], [0, 85, 255], [170, 0, 0], [73, 247, 235]
]

def color_2_label(colors, mod):
    length = colors.shape[0]
    labels = np.zeros([length, 1])
    if mod == 'lower':
        cast = class_2_label_lower
    elif mod == 'upper':
        cast = class_2_label_upper
    for i in range(length):
        color = colors[i]
        for j in range(14):
            if operator.eq(color[0:3].tolist(), cast[j]):
                labels[i] = j + 1
                break
    return labels.astype('int32')

def label_2_color(labels, mod, colors, selected_idx):
    length = labels.shape[0]
    if mod == 'lower':
        cast = class_2_label_lower
    elif mod == 'upper':
        cast = class_2_label_upper
    for i in range(length):
        label = int(labels[i][0])
        if label != 0:
            if selected_idx:
                colors[selected_idx[i]][0] = cast[label - 1][0]
                colors[selected_idx[i]][1] = cast[label - 1][1]
                colors[selected_idx[i]][2] = cast[label - 1][2]
            else:
                colors[i][0] = cast[label - 1][0]
                colors[i][1] = cast[label - 1][1]
                colors[i][2] = cast[label - 1][2]
    return colors

class AverageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def get_avail_gpu():
    result = os.popen("nvidia-smi").readlines()

    try:
        for i in range(len(result)):
            if 'Processes' in result[i]:
                process_idx = i

        num_gpu = 0
        for i in range(process_idx + 1):
            if 'MiB' in result[i]:
                num_gpu += 1
        gpu_list = list(range(num_gpu))

        for i in range(process_idx, len(result)):
            if result[i][22] == 'C':
                gpu_list.remove(int(result[i][5]))
        return (gpu_list[0])
    except:
        print('no gpu available, return 0')
        return 0